#!/usr/bin/env python
#coding=utf-8

import tornado.ioloop
import tornado.web
import tornado.httpclient
import tornado.autoreload
import os, json, time
import requests, re
from pymongo import MongoClient


DBHOST = '133.130.102.237'
DBPORT = 3024
DBNAME = 'sqlscan'

conn = MongoClient(DBHOST, DBPORT)
db = conn[DBNAME]

class BaseHandler(tornado.web.RequestHandler):

	def get_current_user(self):
		return self.get_secure_cookie('u')

# 主页
class MainHandler(BaseHandler):
	@tornado.web.authenticated
	def get(self):
		self.redirect('/tasks')

	def post(self):
		pass
#注销

class ExitHandler(BaseHandler):
	def get(self):
		self.clear_cookie('u')
		self.redirect('/login')
		



# 登录
class LoginHandler(BaseHandler):
	def get(self):

		self.render('login.html', title="login")

	def post(self):
		self.username = self.get_argument('username')
		self.password = self.get_argument('password')

		next = self.get_argument('next', '/')
		a= db.user.find_one({'username':self.username,'password':self.password})
		if a:
			self.write('success')
			self.set_secure_cookie('u', self.username, httponly = True) 
		else:
			self.write('failed')
			
			
			
			
			
class RegHandler(BaseHandler):
	def get(self):
		self.render('register.html',title="注册用户")
	def post(self):
		self.username=self.get_argument('username','')
		self.password=self.get_argument('password','')
		self.email=self.get_argument('email','')
		user=db.user.find_one({'username':self.username})
		if not user:
			db.user.insert({'username':self.username,'password':self.password,'email':self.email})
			self.write('0')
		else:
			self.write('1')

			
# 任务列表
class TasksHandler(BaseHandler):
	@tornado.web.authenticated
	def get(self):
		v = self.get_argument('v', 'html')
		tasks = db.resultlist.find({'username':self.get_current_user()})

		if v == 'html':
			self.render('result.html', tasks = tasks)
		else:
			j = []
			for task in tasks:
				j.append(task)

			self.write(json.dumps(j))

# 任务删除
class TaskDelHandler(BaseHandler):
	@tornado.web.authenticated
	def get(self):
		pass

	@tornado.web.authenticated
	def post(self):
		user=self.get_current_user()
		taskid = self.get_argument('taskid', '')
		test=yield db.resultlist.find_one({'taskid':taskid,'username':user})
		if taskid != '' and not test:
			tasks = db.resultlist.remove({"taskid": taskid})
			self.write('success')
		else:
			self.write('error')

# UserAgent等配置
class ConfigHandler(BaseHandler):
	@tornado.web.authenticated
	def get(self):

		whitelist = db.whitelist.find({'username':self.get_current_user()})
		self.render('config.html', whites = whitelist)

	@tornado.web.authenticated
	def post(self):
		pass



# 白名单添加
class ConfigWhiteAddHandler(BaseHandler):

	@tornado.web.authenticated
	def post(self):
		user=self.get_current_user()
		host = self.get_argument('host', '')
		if '.' in host:

			db.whitelist.insert({"host": host,"username":user})
			self.write('success')
		else:
			self.write('failed')

# 白名单删除
class ConfigWhiteDelHandler(BaseHandler):
	@tornado.web.authenticated
	def post(self):
		
		host = self.get_argument('host', '')
		if host != '':
			db.whitelist.remove({"host": host,'username':self.get_current_user()})
			self.write('success')


class ErrorHandler(tornado.web.RequestHandler):
	def get(self):
		self.write_error(404)

	def write_error(self, status_code, **kwargs):
		if status_code == 404:
			self.set_status(404)
			self.write('404')
		elif status_code == 500:
			self.write('500')
			self.set_status(500)
		else:
			self.write(str(status_code))

#插入任务接口
class TaskInsertAPI(BaseHandler):
        def get(self):
                pass
        def post(self):
            taskjson = self.get_argument('taskjson', '')
            if taskjson != '':
                self.write(taskjson)
                db.tasklist.insert(json.loads(taskjson))
                self.write('success')
#显示白名单内容接口
class WhiteListShwoAPI(BaseHandler):
        def get(self):
                self.post()
        def post(self):
                json1='['
                whitelist=db.whitelist.find({'username':self.get_current_user()})
                for white in whitelist:
                    white=white['host']
                    json1+='{"host":'+'"'+str(white)+'"'+'},'
                json1=json1[:len(json1)-1]
                json1+=']'
                self.write(json1)
class Bugs(tornado.web.RequestHandler):

	@tornado.web.asynchronous
	def get(self):
		ua = 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/45.0.2454.101 Safari/537.36'
		headers = {'User-Agent': ua}
		url = 'http://api.wooyun.org/bugs'
		http = tornado.httpclient.AsyncHTTPClient()
		http.fetch(url, method = 'GET', headers = headers, callback = self.on_response)
        def post(self):
            pass

	def on_response(self, response):
		if response.error:
			self.write('error')
		else:
			self.write(response.body)
		self.finish()


settings = {
	'login_url': '/login',
	'cookie_secret': 'dasdas23',
	'static_path': os.path.join(os.path.dirname(__file__), 'static'),
	'template_path': os.path.join(os.path.dirname(__file__), 'templates'),
	'debug': False,
}

application = tornado.web.Application([
	(r'/', MainHandler),
	(r'/login', LoginHandler),
	(r'/tasks', TasksHandler),
    (r'/taskadd',TaskInsertAPI),
    (r'/whitelistapi',WhiteListShwoAPI),
	# (r'/tasks/([a-zA-Z0-9]+)',TasksViewHandler),
	(r'/tasks/del', TaskDelHandler),
	(r'/config', ConfigHandler),
	(r'/config/white/add', ConfigWhiteAddHandler),
	(r'/config/white/del', ConfigWhiteDelHandler),
	(r'/register',RegHandler),
        (r'/exit',ExitHandler),
	(r'.*', ErrorHandler)
], **settings)
if __name__ == '__main__':
    application.listen(7777)
    tornado.ioloop.IOLoop.instance().start()
